import torch
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
import util
from transformers import TextIteratorStreamer
from threading import Thread

torch.manual_seed(1234)
# Initialize the tokenizer
tokenizer = AutoTokenizer.from_pretrained("hf-models/Qwen-VL", trust_remote_code=True)
# Pass the default decoding hyperparameters of Qwen2-7B-Instruct
# max_tokens is for the maximum length for generation.
# Input the model name or path. Can be GPTQ or AWQ models.
model = AutoModelForCausalLM.from_pretrained("hf-models/Qwen-VL", device_map="cuda", trust_remote_code=True).eval()

# streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
# generation_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=512)


def recognize(image_url, target_text):
    util.download_and_show_image(image_url)

    # Prepare your prompts
    prompt = f'''
    你现在的任务是从图片中提取相关信息，提取目标包括：{target_text}，各提取目标使用|分隔，请自行拆分，不要随意识别不相关信息。
    要求：
    1. 返回结果使用json格式，每条数据包含一个key-value对，key值为我指定的关键信息，value值为所抽取的结果。
    2. 仅考虑单条结果情况。
    3. 如果认为OCR识别结果中没有关键信息key，则将value赋值为“未找到相关信息”即可。如果图片模糊，请不要随意猜测内容，将value赋值为“无法识别”即可。 
    4. 请只输出json格式的结果，不要包含其它多余文字。
    '''
    query = tokenizer.from_list_format([
        {'image': util.image_path},
        {'text': prompt},
    ])
    inputs = tokenizer(query, return_tensors='pt')
    inputs = inputs.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=512)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    generated_text = ""
    for new_text in streamer:
        generated_text += new_text
    print(generated_text)

    # pred = model.generate(**inputs)
    # response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
    # print(response)

    return generated_text